import pandas as pd
import lightgbm as lgb
from datetime import datetime, timedelta
from sklearn.metrics import mean_squared_error

#读取数据集
train_df = pd.read_csv('train_data.csv')
test_df = pd.read_csv('test_data.csv')

#数据处理
train_df['order_date'] = pd.to_datetime(train_df['order_date'])
test_df['order_date'] = pd.to_datetime(test_df['order_date'])

#特征工程
train_df['year'] = train_df['order_date'].dt.year
train_df['month'] = train_df['order_date'].dt.month
train_df['day'] = train_df['order_date'].dt.day
train_df['weekday'] = train_df['order_date'].dt.weekday
test_df['year'] = test_df['order_date'].dt.year
test_df['month'] = test_df['order_date'].dt.month
test_df['day'] = test_df['order_date'].dt.day
test_df['weekday'] = test_df['order_date'].dt.weekday

#定义函数，生成特征
def generate_features(df):
    df['total_price'] = df['item_price'] * df['ord_qty']
    group = df.groupby(['sales_region_code', 'item_code',
                        'first_cate_code', 'second_cate_code',
                        'year', 'month'])['total_price'].sum().reset_index()
    group = group.rename(columns={'total_price': 'target'})
    # 新增一个二分类标签，表示需求量是否大于0
    group['label'] = (group['target'] > 0).astype(int)
    return group

#生成特征
train_group = generate_features(train_df)
test_group = generate_features(test_df)

#准备训练数据和测试数据
train_x = train_group.drop(['target', 'label'], axis=1)
train_y = train_group['label']
test_x = test_group.drop(['target', 'label'], axis=1)

#定义交叉验证函数
def cv(train_x, train_y, params):
    lgtrain = lgb.Dataset(train_x, label=train_y)
    cv_results = lgb.cv(params, lgtrain, num_boost_round=1000, nfold=5, early_stopping_rounds=50, seed=42, verbose_eval=100)
    return cv_results['rmse'][-1]

#定义搜索函数
def objective(params):
    # 将n_estimators转换为整数
    params['n_estimators'] = int(params['n_estimators'])
    params['num_leaves'] = int(params['num_leaves'])
    params['max_depth'] = int(params['max_depth'])
    params['min_data_in_leaf'] = int(params['min_data_in_leaf'])
    params['bagging_freq'] = int(params['bagging_freq'])
    rmse = cv(train_x, train_y, params)
    return {'loss': rmse, 'status': 'ok'}

#搜索最佳参数
from hyperopt import fmin, tpe, hp
space = {
    'max_depth': hp.quniform('max_depth', 5, 15, 1),
    'learning_rate': hp.loguniform('learning_rate', -5, 0),
    'n_estimators': hp.quniform('n_estimators', 50, 500, 10),
    'num_leaves': hp.quniform('num_leaves', 2, 50, 1),
    'min_data_in_leaf': hp.quniform('min_data_in_leaf', 1, 20, 1),
    'feature_fraction': hp.uniform('feature_fraction', 0.1, 1),
    'bagging_fraction': hp.uniform('bagging_fraction', 0.1, 1),
    'bagging_freq': hp.quniform('bagging_freq', 1, 10, 1),
    'lambda_l1': hp.uniform('lambda_l1', 0, 5),
    'lambda_l2': hp.uniform('lambda_l2', 0, 5),
    'random_state': 42
}
best = fmin(fn=objective, space=space, algo=tpe.suggest, max_evals=1000)

#使用最佳参数训练模型
params = {
    'max_depth': int(best['max_depth']),
    'learning_rate': float(best['learning_rate']),
    'n_estimators': int(best['n_estimators']),
    'num_leaves': float(best['num_leaves']),
    'min_data_in_leaf': float(best['min_data_in_leaf']),
    'feature_fraction': float(best['feature_fraction']),
    'bagging_fraction': float(best['bagging_fraction']),
    'bagging_freq': int(best['bagging_freq']),
    'lambda_l1': float(best['lambda_l1']),
    'lambda_l2': float(best['lambda_l2']),
    'random_state': 42
}
model = lgb.LGBMClassifier(**params)
model.fit(train_x, train_y)

#预测测试集的标签
test_group['label'] = model.predict(test_x)

#保存结果
result_df = test_group[['sales_region_code', 'item_code', 'first_cate_code', 'second_cate_code',
            'year', 'month', 'label']].rename(columns={'label': 'demand_flg'})
result_df.to_csv('result.csv', index=False)